# -*- coding: utf-8 -*-
# @Author : Gan
# @Time : 2024/9/12 21:06

import os
import cv2
import numpy as np
import random
import shutil
from glob import glob
from tqdm import tqdm
from osgeo import gdal, osr

def add_gaussian_noise(image, mean=0, var=1):
    """
    向图像添加高斯噪声
    :param image: OpenCV图像对象（numpy数组）
    :param mean: 噪声的均值
    :param var: 噪声的方差
    :return: 带有噪声的OpenCV图像对象（numpy数组）
    """
    row, col, ch = image.shape
    gauss = np.random.normal(mean, var ** 0.5, (row, col, ch))
    gauss = gauss.reshape(row, col, ch)
    noisy = image + gauss

    # 确保像素值不会超出0到255的范围
    noisy = np.clip(noisy, 0, 255).astype(np.uint8)
    return noisy


def rotate_image(img, angle):
    (height, width) = img.shape[:2]

    # 计算旋转中心（图像中心）
    center = (width // 2, height // 2)
    # 计算旋转矩阵
    rotation_matrix = cv2.getRotationMatrix2D(center, angle, scale=1.0)

    # 计算旋转后图像的新边界尺寸，确保图像完整
    cos = np.abs(rotation_matrix[0, 0])
    sin = np.abs(rotation_matrix[0, 1])
    new_width = int((height * sin) + (width * cos))
    new_height = int((height * cos) + (width * sin))

    # 调整旋转矩阵，使旋转后的图像中心与原图中心对齐
    rotation_matrix[0, 2] += (new_width / 2) - center[0]
    rotation_matrix[1, 2] += (new_height / 2) - center[1]

    # 执行旋转并获取结果
    rotated_img = cv2.warpAffine(img, rotation_matrix, (new_width, new_height))

    return rotated_img


def get_crop_size_list(mu, sigma, num, min_size, max_size):
    random_numbers = np.random.normal(mu, sigma, num)
    random_size = list(np.clip(random_numbers, min_size, max_size))
    return random_size


def crop_background(img):
    alpha_channel = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    _, non_transparent = cv2.threshold(alpha_channel, 0, 255, cv2.THRESH_BINARY)
    # 获取非透明区域的边界框
    contours, _ = cv2.findContours(non_transparent, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    x, y, w, h = cv2.boundingRect(contours[0])
    return x,x+w,y,y+h
    # 裁剪图像
    cropped_image = img[y:y + h, x:x + w]
    return cropped_image
def pixel_to_geo_coords(geo_transform, px, py):
    """从像素坐标转换到地理坐标"""
    ulx, xres, xskew, uly, yskew, yres = geo_transform
    lon = ulx + (px + 1) * xres + py * xskew
    lat = uly + (px + 1) * yskew + py * yres
    return lon, lat
def cut_image(src_ds, x_min, y_min, x_max, y_max,output_file=''):


    # 获取图像的地理变换矩阵
    geo_transform = src_ds.GetGeoTransform()
    if geo_transform is None:
        print("Failed to get GeoTransform from input file.")
        return

    # # 计算裁剪区域在像素坐标系下的左上角和右下角位置
    # x_min = int((ulx - geo_transform[0]) / geo_transform[1])
    # y_max = int((uly - geo_transform[3]) / geo_transform[5])
    # x_max = int((lrx - geo_transform[0]) / geo_transform[1])
    # y_min = int((lry - geo_transform[3]) / geo_transform[5])

    # 裁剪区域宽度和高度
    width = x_max - x_min
    height = y_max - y_min
    lon, lat = pixel_to_geo_coords(geo_transform,x_min, y_min)
    # 创建一个新的GeoTransform矩阵
    new_geo_transform = list(geo_transform)
    new_geo_transform[0] = lon
    new_geo_transform[3] = lat

    # 创建一个内存中的临时数据集用于裁剪
    mem_drv = gdal.GetDriverByName('MEM')
    dst_ds = mem_drv.Create('', width, height, src_ds.RasterCount, src_ds.GetRasterBand(1).DataType)
    if dst_ds is None:
        print('Failed to create memory raster.')
        return

    # 设置新的地理变换矩阵
    dst_ds.SetGeoTransform(new_geo_transform)

    # 复制投影信息
    dst_ds.SetProjection(src_ds.GetProjection())

    # 执行裁剪
    # 注意: 这里使用src_ds.ReadAsArray方法直接从源数据集中读取指定范围的数据
    data = src_ds.ReadAsArray(x_min, y_min, width, height)

    for band_num in range(1, src_ds.RasterCount + 1):
        band = dst_ds.GetRasterBand(band_num)
        band.WriteArray(data[band_num - 1])
    del src_ds
    if output_file:
    # 将内存中的图像保存到磁盘
        save_tif(output_file,dst_ds)
    return dst_ds




def crop_img(img, crop_size, thresh):
    w = img.RasterXSize
    h = img.RasterYSize
    # h, w = img.shape[:2]

    crop_h = int(crop_size)
    crop_w = int(crop_size * random.choice(scale_factor))

    coord_h, coord_w = random.randint(0, h - thresh), random.randint(0, w - thresh)
    end_h, end_w = min(coord_h + crop_h, h), min(coord_w + crop_w, w)
    croped_img = cut_image(img,coord_w, coord_h, end_w, end_h)
    # croped_img = img[coord_h:end_h, coord_w:end_w, :]

    return croped_img, coord_h, coord_w, end_h, end_w
def transform_point_from_wgs84_to_web_mercator(latitude, longitude):
    """将WGS 84坐标系下的点转换到Web Mercator坐标系"""
    # 创建WGS 84坐标系
    wgs84_crs = osr.SpatialReference()
    wgs84_crs.ImportFromEPSG(4326)

    # 创建Web Mercator坐标系
    web_mercator_crs = osr.SpatialReference()
    web_mercator_crs.ImportFromEPSG(3857)

    # 创建坐标变换对象
    coord_transform = osr.CoordinateTransformation(wgs84_crs, web_mercator_crs)

    # 将经纬度坐标转换为Web Mercator坐标
    web_mercator_x, web_mercator_y, _ = coord_transform.TransformPoint(longitude, latitude)

    return web_mercator_x, web_mercator_y
def geo2geo(ds_a,ds_b,lon,lat):
    '''
    将B影像中点的地理坐标转换为A中的地理坐标
    Args:
        ds_a: A影像(gdal)
        ds_b: A影像(gdal)
        lon:B中点的纬度
        lat:B中点的经度

    Returns:点在A中经纬度

    '''
    crs_a = osr.SpatialReference(wkt=ds_a.GetProjection())
    crs_b = osr.SpatialReference(wkt=ds_b.GetProjection())
    # 将 B 图像的左上角经纬度转换为 A 图像的 CRS 中的坐标
    if not crs_b.IsSame(crs_a):
        ct = osr.CoordinateTransformation(crs_b, crs_a.CloneGeogCS())
        upper_left_lat, upper_left_lon, _ = ct.TransformPoint(lat, lon)
        x,y = transform_point_from_wgs84_to_web_mercator(upper_left_lat, upper_left_lon)
        return y,x
    # return upper_left_lat, upper_left_lon
def geo2pixel(ds,lon,lat):
    '''
    将地理坐标转换为像素坐标
    Args:
        ds: 影像(gdal)
        lon: 经度
        lat: 纬度

    Returns:列号，行号

    '''
    geo_transform_a = ds.GetGeoTransform()
    # 将经纬度转换为 A 图像的像素坐标
    col = int((lon - geo_transform_a[0]) / geo_transform_a[1])
    row = int((lat - geo_transform_a[3]) / geo_transform_a[5])
    return col, row
def pixel_coordinates_of_upper_left_corner2A(a_path, ds_b):
    """
    从 B 图像的地理信息中读取左上角点的经纬度，并根据这些经纬度计算在 A 图像中的像素坐标。

    参数:
    a_path (str): A 图像的文件路径。
    b_path (str): B 图像的文件路径。

    返回:
    tuple: (row, col) 在 A 图像中的像素坐标。
    """
    # 读取 A 图像的变换矩阵和 CRS
    ds_a = gdal.Open(a_path)
    lon, lat = ds_b.GetGeoTransform()[0], ds_b.GetGeoTransform()[3]
    upper_left_lat, upper_left_lon = geo2geo(ds_a,ds_b,lat,lon)
    col, row = geo2pixel(ds_a, upper_left_lon, upper_left_lat)

    return col, row

def pixel_coordinates_of_lower_right_corner2A(a_path, ds_b):
    lon, lat = get_pixel_coordinates_of_lower_right_corner(ds_b)
    ds_a = gdal.Open(a_path)
    upper_right_lat, upper_right_lon = geo2geo(ds_a, ds_b, lat, lon)
    col, row = geo2pixel(ds_a, upper_right_lon, upper_right_lat)
    return col, row

def get_pixel_coordinates_of_lower_right_corner(ds_b):

    # 获取图像的宽度和高度
    width = ds_b.RasterXSize
    height = ds_b.RasterYSize
    # 获取地理变换矩阵
    geo_transform = ds_b.GetGeoTransform()

    # 计算右下角点的地理坐标
    ulx = geo_transform[0]  # 左上角X坐标
    uly = geo_transform[3]  # 左上角Y坐标
    xres = geo_transform[1]  # 沿X轴方向每个像素的宽度
    yres = geo_transform[5]  # 沿Y轴方向每个像素的高度
    lrx = ulx + width * xres + height * geo_transform[2]
    lry = uly + width * geo_transform[4] + height * yres

    # 返回图像的宽度、高度和右下角点的地理坐标
    return lrx, lry

def read_tif(input_file):
    # 打开原始图像
    src_ds = gdal.Open(input_file, gdal.GA_ReadOnly)
    if src_ds is None:
        print('Unable to open', input_file)
    return src_ds
def gdal_to_opencv(src_ds):
    # 打开TIFF图像
    # src_ds = gdal.Open(input_file, gdal.GA_ReadOnly)
    # if src_ds is None:
    #     print('无法打开文件:', input_file)
    #     return None

    # 获取图像的宽度和高度
    width = src_ds.RasterXSize
    height = src_ds.RasterYSize
    bands = src_ds.RasterCount

    # 根据波段数决定图像的颜色模式
    if bands == 1:  # 单波段图像（灰度）
        band = src_ds.GetRasterBand(1)
        data = band.ReadAsArray(0, 0, width, height)
        img = np.array(data, dtype=np.uint8)
    elif bands == 3:  # 三波段图像（BGR）
        b_band = src_ds.GetRasterBand(1)
        g_band = src_ds.GetRasterBand(2)
        r_band = src_ds.GetRasterBand(3)

        b = b_band.ReadAsArray(0, 0, width, height)
        g = g_band.ReadAsArray(0, 0, width, height)
        r = r_band.ReadAsArray(0, 0, width, height)

        img = np.dstack((b, g, r))  # 注意这里的顺序
        img = np.array(img, dtype=np.uint8)
        return img
    elif bands == 4:  # 四波段图像（通常为BGR+NIR或BGR+Alpha）
        b_band = src_ds.GetRasterBand(1)
        g_band = src_ds.GetRasterBand(2)
        r_band = src_ds.GetRasterBand(3)
        nir_band = src_ds.GetRasterBand(4)  # 假设第四个波段是NIR

        b = b_band.ReadAsArray(0, 0, width, height)
        g = g_band.ReadAsArray(0, 0, width, height)
        r = r_band.ReadAsArray(0, 0, width, height)
        nir = nir_band.ReadAsArray(0, 0, width, height)

        # 创建BGR图像
        img_bgr = np.dstack((b, g, r))  # 注意这里的顺序
        img_bgr = np.array(img_bgr, dtype=np.uint8)

        # 如果需要，也可以处理NIR波段
        # img_nir = np.array(nir, dtype=np.uint8)

        # 返回BGR图像
        return img_bgr
    else:
        print("不支持的波段数:", bands)
        return None
def save_tif(output_file,dst_ds):
    if output_file:
    # 将内存中的图像保存到磁盘
        out_driver = gdal.GetDriverByName('GTiff')
        out_data = out_driver.CreateCopy(output_file, dst_ds, 0)
        return out_data
    else:
        print('output_file is False')

# 通过缩放控制h和w的尺寸差异
scale_factor = [0.8, 0.9, 1, 1.1, 1.2]
rotate_angle_list = list(np.arange(-50, 51, 10))

if __name__ == '__main__':
    # 随机裁剪生成查询图像，只针对波总下载的数据，需要将路径中所有中文手动删除
    # root_path = r'H:\新建文件夹'
    root_path = 'G:/data/satgeoloc_finals_basemap/query'
    raster_dst_path = 'G:/data/satgeoloc_finals_basemap/query_image/raster'
    tif_dst_path = 'G:/data/satgeoloc_finals_basemap/query_image/tif'
    basemap_dir = 'G:/data/satgeoloc_finals_basemap/google'
    os.makedirs(raster_dst_path, exist_ok=True)
    os.makedirs(tif_dst_path, exist_ok=True)
    img_list = glob(f'{root_path}/*/*/*.tif')
    for img_path in tqdm(img_list):
        name, zoom = img_path.split('\\')[-3], img_path.split('\\')[-2]
        img0 = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
        x1,x2,y1,y2 = crop_background(img0)
        img = read_tif(img_path)
        output_path = ''
        # temp = img_path.split('orgin')
        # output_path = temp[0]+'cut_orgin'+temp[1]
        # os.makedirs(os.path.dirname(output_path),exist_ok=True)
        img = cut_image(img, x1, y1, x2, y2, output_path)
        # img_x = img.RasterXSize
        # img_y = img.RasterYSize
        # if (img_y,img_x)!=img0.shape[0:2]:
        #     cv2.imwrite(img_path, img)
        basemap_name = os.path.basename(os.path.dirname(os.path.dirname(img_path)))+'.tif'
        basemap_path = os.path.join(basemap_dir,basemap_name)

        # img = gdal_to_opencv(img)
        if zoom == '4':
            random_size = get_crop_size_list(3500, 500, 30, 500, 5000)
            thresh = 1000
        elif zoom == '1':
            random_size = get_crop_size_list(5500, 2000, 50, 1000, 10000)
            thresh = 4000
        elif zoom == '8':
            random_size = get_crop_size_list(1800, 300, 50, 1000, 3000)
            thresh = 500
        else:
            coord_h = coord_w = 0
            end_h = img.RasterYSize
            end_w = img.RasterXSize
            x1, y1 = pixel_coordinates_of_upper_left_corner2A(basemap_path, img)
            x2,y2 = pixel_coordinates_of_lower_right_corner2A(basemap_path,img)
            save_tif(f'{tif_dst_path}/{name}_{zoom}_{coord_w}_{coord_h}_{end_w}_{end_h}_{x1}_{y1}_{x2}_{y2}.tif', img)
            img = gdal_to_opencv(img)
            cv2.imwrite(f'{raster_dst_path}/{name}_{zoom}_{coord_w}_{coord_h}_{end_w}_{end_h}_{x1}_{y1}_{x2}_{y2}.tif', img)
            continue

        for crop_size in random_size:
            rotate_angle = random.choice(rotate_angle_list)

            croped_img, coord_h, coord_w, end_h, end_w = crop_img(img, crop_size, thresh)
            x1, y1 = pixel_coordinates_of_upper_left_corner2A(basemap_path, croped_img)
            x2, y2 = pixel_coordinates_of_lower_right_corner2A(basemap_path, croped_img)
            save_tif(f'{tif_dst_path}/{name}_{zoom}_{coord_w}_{coord_h}_{end_w}_{end_h}_{x1}_{y1}_{x2}_{y2}.tif', croped_img)
            croped_img = gdal_to_opencv(croped_img)
            croped_img = add_gaussian_noise(croped_img)
            result = rotate_image(croped_img, rotate_angle)
            cv2.imwrite(f'{raster_dst_path}/{name}_{zoom}_{coord_w}_{coord_h}_{end_w}_{end_h}_{x1}_{y1}_{x2}_{y2}.tif', result)

